#!/usr/bin/env python
# encoding: utf-8

import os
import argparse
import numpy as np
import pandas as pd
from utils import compute_eer

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--trials_path', help='trials path', type=str, default="trials.lst")
    parser.add_argument('--embedding_dim', help='embedding dim', type=int, default=512)
    args = parser.parse_args()

    print("loading trials")
    trials = np.load(args.trials_path)
    trials_df = pd.DataFrame(trials, columns=["enroll", "test", "target"])

    target_matrix = trials_df.pivot(index="enroll", columns="test", values="target")

    cnt = 0
    for item in target_matrix.values:
        for val in item:
            if val == '1':
                cnt += 1
    input(cnt)

    enroll_list = target_matrix.index
    test_list = target_matrix.columns

    print("loading enroll")
    enroll_vectors = np.zeros((len(enroll_list), args.embedding_dim), dtype=np.float)
    for idx, enroll in enumerate(enroll_list):
        enroll_vectors[idx] = np.load(enroll)

    print("loading test")
    test_vectors = np.zeros((len(test_list), args.embedding_dim), dtype=np.float)
    for idx, test in enumerate(test_list):
        test_vectors[idx] = np.load(test)

    print("scoring")
    cosine_matrix = enroll_vectors.dot(test_vectors.T)
    enroll_norm = np.linalg.norm(enroll_vectors, axis=1, keepdims=True)
    test_norm = np.linalg.norm(test_vectors, axis=1, keepdims=True)
    denom = enroll_vectors.dot(test_vectors.T)
    cosine_matrix = cosine_matrix/enroll_norm.dot(test_norm.T)

    target_matrix = target_matrix.values
    target_scores = []
    nontarget_scores = []
    for x in range(len(target_matrix)):
        for y in range(len(target_matrix[0])):
            if target_matrix[x][y] == '1':
                target_scores.append(cosine_matrix[x][y])
            else:
                nontarget_scores.append(cosine_matrix[x][y])

    print("number of target: {}".format(len(target_scores)))
    print("number of nontarget: {}".format(len(nontarget_scores)))

    eer, threshold = compute_eer(target_scores, nontarget_scores)
    print("EER: {}".format(eer))
    print("Threshold: {}".format(threshold))

